package cn.swing.main.srv.cv;

import cn.swing.main.srv.cv.model.TennisPoint;
import cn.swing.main.srv.cv.tracker.AdvancedTracker;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Date;
import java.util.List;
import java.util.stream.Collectors;
import org.opencv.core.*;
import org.opencv.highgui.HighGui;
import org.opencv.imgproc.Imgproc;
import org.opencv.video.BackgroundSubtractorMOG2;
import org.opencv.video.Video;
import org.opencv.videoio.VideoCapture;
import org.opencv.videoio.Videoio;

public class TennisBallTracking extends BaseOpencv {

    //    private static final int HISTORY_SIZE = 5; // 历史中心位置列表的大小
    //
    //    private static final double DISTANCE_THRESHOLD = 20; // 连续运动的距离阈值
    //
    //    private List<Point> centerHistory = new ArrayList<>();

    public static void main(String[] args) {
        List<String> videos = new ArrayList<>();
        videos.add("2291a2d1-5986-4ece-9eb1-3131f8a559a7-1");
        videos.add("2691751a-af5a-4811-b00f-3ede2d141aeb-1");
        videos.add("2a5c8891-a4c2-4bf5-b8b5-e2b80fe6741f-1");
        videos.add("2b136df9-007c-4a99-88b4-f9b366736a0a-1");
        videos.add("2d75fe27-a79f-42f4-aa8b-083f14fc4a04-1");
        videos.add("2df1bd8f-5232-453e-a10b-502bfae2727e-1");
        videos.add("310fcb77-97c4-44cb-afdd-31ea8fedf51c-1");
        videos.add("3395a3cb-ec74-4c96-b39a-413f17b39caa-1");
        videos.add("3f5bfd1d-4d33-4c3d-846a-24dd48855f4c-1");
        videos.add("402af8d5-1888-4d0a-b42d-871344f1665f-1");
        videos.add("40f94123-de3f-48e5-88b6-f396d682d1df-1");
        videos.add("4609ba7f-27cb-45da-8cf3-4c98c5cf36e2-1");
        videos.add("462386cb-3398-40d8-9f19-b0113313f026-1");
        videos.add("4cc4cb12-24d6-422d-b87c-a302ca096fa4-1");
        videos.add("4d5144bf-2e8a-4f3a-acf7-7964bb6b6596-1");
        videos.add("4d83c0a9-52c9-4025-bd28-bf1c654b91fb-1");
        videos.add("51f6bac9-f013-48e8-9f67-08a5edcd5de1-1");
        videos.add("55197dfc-8363-40bb-8a9a-cb338adf8144-1");
        videos.add("5674ccdd-4768-4203-8f89-26ed0ebad510-1");
        videos.add("5981632c-a09b-4d5f-8c58-7b4b557f9e22-1");
        videos.add("5a485ff0-ba28-42f3-90ea-347a343a406f-1");
        videos.add("5ab8d7db-47e0-41c8-ae1d-3ec81a480350-1");
        videos.add("5b514b57-b683-4aac-96e7-3eed972c2353-1");
        videos.add("5b5878b9-b1bf-43d4-b513-bd8ce7f7e40c-1");
        videos.add("5dc0f720-4982-4fa2-98fe-8b6fbc007ce7-1");
        videos.add("5eba2346-d1a7-412d-a1b0-7de2700d0998-1");
        videos.add("5eec2a01-3fcd-46f4-aa07-0b8f9bc56e15-1");
        videos.add("60b589df-ecb4-44ee-8adf-e10443bf3d95-1");
        videos.add("60fe753c-5def-4ac4-88b8-111c0edf210e-1");
        videos.add("654f99de-b041-4cbd-8e73-a098304aede4-1");
        videos.add("66205081-e3bf-4eba-a157-86678117f7e7-1");
        videos.add("6631f888-0472-4635-b2d7-327b14332fbc-1");
        videos.add("6e8a37c4-160f-434d-93fd-03f9622981e8-1");
        videos.add("6e9a21b2-7534-47c2-9f8b-dc4e28c25166-1");
        videos.add("6f018aa7-53a6-4817-92b5-c38ce3935361-1");
        videos.add("703b1e1c-a9b3-4f2a-bdfd-9d1b29703307-1");
        videos.add("71863583-b657-4635-aab0-8ed7755a2f56-1");
        videos.add("7251ece2-0735-4c3a-8656-5bf5c4a7cb8d-1");
        videos.add("743f55f4-9570-424f-afa6-dee972d90e11-1");
        videos.add("75ad3c92-016f-48eb-b4b2-13908a5bc83d-1");
        videos.add("796f629e-c85e-4218-84a8-620877e9720a-1");
        videos.add("79da5e0c-6d19-4b2e-b62f-b1466c04ec8f-1");
        videos.add("7c20c8de-136c-47fa-b1a0-a885215e61a6-1");
        videos.add("7dc31551-6912-4e3e-a438-c8740ea9c91f-1");
        videos.add("7e475792-d9cf-4109-a8d9-fe46133a1c27-1");
        videos.add("7f91d8e9-6d0d-46c7-8a25-f0f013149a84-1");
        videos.add("80bba770-ca50-47f9-904b-d52cbecc3658-1");
        videos.add("841c0f63-a049-4eed-ad0a-c08c46f82aee-1");
        videos.add("9138ae46-6494-42b7-85ad-13052ba4d056-1");
        videos.add("91c9e40e-66ab-42b5-8135-972a3923db17-1");
        videos.add("924af4b6-2c1b-42eb-9da5-5f4d58d415f8-1");
        videos.add("935915e5-0ca8-4518-a419-eac52050cf95-1");
        videos.add("93ea71e0-60f2-4d24-a8c1-45c18b75c813-1");
        videos.add("951d25ed-e3bd-4472-8cb0-630e58cb8c70-1");
        videos.add("9521e250-96a9-47ba-8ad4-150b333a1a26-1");
        videos.add("9fc1d531-b3d5-4515-88e3-fd1557ebc861-1");
        videos.add("a39c7275-d4d0-4146-9e2f-ec564d3f33d6-1");
        videos.add("a6be094d-1b65-460e-ad4e-8301ac148c61-1");
        videos.add("a77c16a0-2098-428e-8b38-9073ed4032c5-1");
        videos.add("a7fcf676-95d4-4b78-9a2c-b945bbd6f9d7-1");
        videos.add("a8bd3ae4-3a97-4962-a9bf-67fc528f262c-1");
        videos.add("aa3fd4e9-7583-4143-8410-178af0886e3c-1");
        videos.add("ab137990-b14a-4523-8bd3-8f6cc3942496-1");
        videos.add("affd8789-42ac-4e7a-87d1-17453435546f-1");
        videos.add("bb1575b8-6e3d-4a0c-98eb-a4a17345b655-1");
        videos.add("bf08bde0-b8ea-478e-84ce-0cbabd8f1ae3-1");
        videos.add("c5d51b1e-6fb0-4894-ae2c-ebf97bf263f3-1");
        videos.add("c961ddbf-292b-4d6d-ac28-d66e1d94ad60-1");
        videos.add("cb3dd23e-9482-42e4-b7e9-9f183576ff8e-1");
        videos.add("ccc1735c-0e24-44e2-bc92-ff6727d1e524-1");
        videos.add("cdfb1c9a-702c-450a-8c5d-abab64549ac0-1");
        videos.add("cecac45e-19ed-4616-b0cd-ec5fd41d0c8a-1");
        videos.add("d4169546-e15b-499e-a70c-8a2067894588-1");
        videos.add("d42ea042-4a29-4a5f-9635-91c9d63f6d44-1");
        videos.add("d7295fb9-455a-4f69-b804-11466648f1d1-1");
        videos.add("d9881b25-8122-4564-b0c2-3ae53010006a-1");
        videos.add("e042e3ee-f7eb-46f4-ae0f-046679667a6e-1");
        videos.add("e0c1cdee-f447-434e-be04-e91606e1925f-1");
        videos.add("e4d4290e-a840-499e-8d31-0059bf2e868b-1");
        videos.add("e883498a-e996-4572-a0b9-6047ed57a8b3-1");
        videos.add("ec4e374a-2f46-4c64-8914-40b93db9b869-1");
        videos.add("f0f69523-f757-488f-a092-4c126e1f3511-1");
        videos.add("f137a983-4de4-4823-a702-070b1e403f83-1");
        videos.add("f1711ba1-bfe9-4e10-b64b-8e5762bcbed5-1");
        videos.add("f3af73d7-9764-46d7-b16a-934e77af98b0-1");
        videos.add("f44a8f46-26af-4c9b-aa5f-a98161742a7f-1");
        videos.add("fcb03746-b8ea-4df7-8c47-564ba6cbbf2e-1");
        videos.add("ffb34c2b-3364-4887-8acb-d76ec1202507-1");

        /**
         * fabaa955-26d4-4020-ac16-c3a7035fbda5-1__53400_54000 红色树胶场地，背景树丛，白天，阳光充足 亮度平均132 top亮度123
         * 5a485ff0-ba28-42f3-90ea-347a343a406f-1  蓝色树胶场地，背后黄红色的房子，白天，黄昏，阳光充足 亮度平均116 top亮度172
         * 3f5bfd1d-4d33-4c3d-846a-24dd48855f4c-1  白天，阳光，蓝色树胶场地，背后有一栋白色房子 树木 亮度平均112 top亮度109
         * 6e9a21b2-7534-47c2-9f8b-dc4e28c25166-1  晚上户外 前面很晃 亮度平均126 top亮度30
         * 6f018aa7-53a6-4817-92b5-c38ce3935361-1  视角低 草地 双打  亮度平均145 top亮度205
         *
         * @param videoId
         */
        // videos.add("6f018aa7-53a6-4817-92b5-c38ce3935361-1");
        // 开始秒
        int startSecond = 40;
        // 持续时间
        int durationSecond = 30;
        // 调试球追踪
        boolean debugBallTracking = false;
        // 调试轨迹
        boolean debugTrajectory = false;
        // 输出视频
        boolean outputFlag = true;

        videos.forEach(videoId -> new TennisBallTracking().run(startSecond, durationSecond,
                videoId, debugBallTracking, debugTrajectory, outputFlag));

    }

    public void run(int startSecond, int durationSecond, String videoId, boolean debugBallTracking,
            boolean debugTrajectory, boolean outputFlag) {
        String path = "D:\\swing\\" + videoId + ".mp4";
        VideoCapture capture = new VideoCapture(path);
        if (!capture.isOpened()) {
            System.out.println("Error: Unable to open video file.");
            return;
        }

        // 创建背景减除器
        BackgroundSubtractorMOG2 subtractor = Video.createBackgroundSubtractorMOG2();
        subtractor.setHistory(60);
        subtractor.setBackgroundRatio(0.7);
        subtractor.setNMixtures(5);
        subtractor.setVarThreshold(25);
        subtractor.setDetectShadows(false);

        // 获取视频时长
        final double totalFrameCount = capture.get(Videoio.CAP_PROP_FRAME_COUNT);
        // 获取视频的宽度和高度
        int frameWidth = (int) capture.get(Videoio.CAP_PROP_FRAME_WIDTH);
        int frameHeight = (int) capture.get(Videoio.CAP_PROP_FRAME_HEIGHT);
        double fps = capture.get(Videoio.CAP_PROP_FPS);
        double startFrame = startSecond * fps;
        double stopFrame = startFrame + durationSecond * fps;

        int frameCount = 0;
        // 追踪器
        AdvancedTracker tracker = new AdvancedTracker();

        long now = System.currentTimeMillis();
        double totalDealTime = Math.min(durationSecond, Math.floor(totalFrameCount / fps));
        System.out.println(new Date() + ", 即将处理视频" + videoId + ", 总时长(s)=" + totalDealTime);
        long traceTime = 0;
        while (true) {
            Mat frame = new Mat();
            Mat fgMask = new Mat();
            Mat grayFrame = new Mat();
            Mat morphMask = new Mat();
            try {
                if (++frameCount % 60 == 0) {
                    System.out.println(new Date() + ", 处理进度: " + (frameCount / fps) + "s / " + totalDealTime + "s");
                }
                if (!capture.read(frame)) {
                    break;
                }
                if (frameCount > stopFrame) {
                    break;
                }
                if (frameCount < startFrame) {
                    continue;
                }
                // 灰度化
                Imgproc.cvtColor(frame, grayFrame, Imgproc.COLOR_BGR2GRAY);
                // 计算上部分1/5图像的平均亮度 (一般是天空，如果是室外晚上，会非常低，)
                Rect top = new Rect(0, 0, frameWidth, frameHeight / 5);
                Scalar averageMean = Core.mean(grayFrame);
                Scalar topMean = Core.mean(grayFrame.submat(top));
                // 根据平均亮度判断光线条件
                double averageBrightness = averageMean.val[0];
                double topBrightness = topMean.val[0];

                // 背景减除
                subtractor.apply(grayFrame, fgMask);

                // 形态学操作：开运算（去除噪声）
                Mat kernel = Imgproc.getStructuringElement(Imgproc.MORPH_RECT, new Size(3, 3));
                Imgproc.morphologyEx(fgMask, morphMask, Imgproc.MORPH_OPEN, kernel);

                // 形态学操作：闭运算（连接断裂的区域）
                Imgproc.morphologyEx(morphMask, morphMask, Imgproc.MORPH_CLOSE, kernel);
                kernel.release();
                // 检测轮廓
                List<MatOfPoint> contours = new ArrayList<>();
                Imgproc.findContours(morphMask.clone(), contours, new Mat(), Imgproc.RETR_EXTERNAL,
                        Imgproc.CHAIN_APPROX_SIMPLE);

                List<TennisPoint> points = new ArrayList<>();
                for (MatOfPoint contour : contours) {
                    // 计算轮廓的面积
                    double area = Imgproc.contourArea(contour);
                    // 计算轮廓的周长
                    double perimeter = Imgproc.arcLength(new MatOfPoint2f(contour.toArray()), true);
                    // 计算圆形度
                    double circularity = (4 * Math.PI * area) / (perimeter * perimeter);

                    Rect boundingBox = Imgproc.boundingRect(contour);
                    // 颜色亮度特征计算
                    TennisPoint tennisPoint = caculationColor(boundingBox, frameCount, frame, averageBrightness,
                            topBrightness, circularity);
                    // 过滤面积和区域
                    tennisPoint = filterByColorAndArea(tennisPoint);
                    if (tennisPoint != null) {
                        points.add(tennisPoint);
                    }
                }
                // 显示处理后的帧 绘制追踪框
                if (debugBallTracking) {
                    for (int i = 0; i < points.size(); i++) {
                        TennisPoint point = points.get(i);
                        Scalar text;
                        if (point.getFilterType() == 0) {
                            text = new Scalar(0, 0, 180 + Math.abs(i % 75));
                        } else {
                            text = new Scalar(255, 0, 0);
                        }
                        Imgproc.rectangle(frame, new Point(point.getX() - point.getWidth() / 2,
                                point.getY() - point.getHeight() / 2), new Point(point.getX() + point.getWidth() / 2,
                                point.getY() + point.getHeight() / 2), new Scalar(255, 0, 0), 3);

                        // 在框框的右边绘制轨迹id
                        Imgproc.putText(frame, i + "",
                                new Point(point.getX() + point.getWidth(), point.getY()),
                                Imgproc.FONT_HERSHEY_SIMPLEX, 0.6,
                                text, 2);
                    }
                    HighGui.imshow("Tennis Ball Tracking", frame);
                    final int i = HighGui.waitKey(30);
                }
                // 处理帧识别点
                traceTime += tracker.processFrame(points.stream()
                        .filter(p -> p.getFilterType() == 0)
                        .collect(Collectors.toList()), frameCount);

            } finally {
                frame.release();
                fgMask.release();
                grayFrame.release();
                morphMask.release();
            }
        }
        capture.release();
        tracker.flushAndMerge();
        tracker.debugShow(debugTrajectory, outputFlag, path, startFrame, stopFrame);
        HighGui.destroyAllWindows();
        System.out.println("视频处理耗时:" + (System.currentTimeMillis() - now));
        System.out.println("点数据追踪耗时:" + traceTime);
        System.out.println("结束，总耗时:" + (System.currentTimeMillis() - now));
    }

    private TennisPoint caculationColor(Rect boundingBox, int frameCount, Mat frame, double averageBrightness,
            double topBrightness, double circularity) {
        double centerX = boundingBox.x + boundingBox.width / 2.0;
        double centerY = boundingBox.y + boundingBox.height / 2.0;
        // 获取frame在当前rect区域的颜色信息 获取中心点1/2的区域
        Rect colorBox = new Rect((int) centerX - boundingBox.width / 4, (int) centerY - boundingBox.height / 4,
                boundingBox.width / 2, boundingBox.height / 2);
        Mat ballColor = frame.submat(colorBox);
        // 创建一个1x1的Mat来存储HSV颜色
        Mat hsvMat = new Mat(1, 1, CvType.CV_8UC3);

        // 将BGR颜色转换为HSV颜色
        Imgproc.cvtColor(ballColor, hsvMat, Imgproc.COLOR_BGR2HSV);
        // 提取HSV颜色值
        double[] hsvValues = hsvMat.get(0, 0);

        // 返回HSV颜色的Scalar对象
        Scalar hsvColor = new Scalar(hsvValues[0], hsvValues[1], hsvValues[2]);

        // 取ballColor的颜色的中位数 BGR
        Scalar medianColor = getMedianColor(ballColor);
        Scalar meanColor = Core.mean(ballColor);
        Mat ballBright = new Mat();
        Imgproc.cvtColor(ballColor, ballBright, Imgproc.COLOR_BGR2GRAY);
        Scalar averageMean = Core.mean(ballBright);
        double bright = averageMean.val[0];

        final TennisPoint tennisPoint = new TennisPoint(centerX, centerY, boundingBox.width, boundingBox.height,
                frameCount);
        tennisPoint.setCircularity(Math.ceil(circularity * 100) / 100);
        tennisPoint.setBright(Math.ceil(bright * 100) / 100);
        tennisPoint.setFrameBright(Math.ceil(averageBrightness * 100) / 100);
        tennisPoint.setFrameTopBright(Math.ceil(topBrightness * 100) / 100);
        tennisPoint.setMedianColor(medianColor);
        tennisPoint.setMeanColor(meanColor);
        tennisPoint.setHsvColor(hsvColor);
        return tennisPoint;
    }

    public Scalar getMedianColor(Mat ballColor) {
        // 获取子矩阵的宽度和高度
        int width = ballColor.width();
        int height = ballColor.height();

        // 存储每个通道的像素值
        double[] blueValues = new double[width * height];
        double[] greenValues = new double[width * height];
        double[] redValues = new double[width * height];

        // 提取每个像素的BGR值
        int index = 0;
        for (int i = 0; i < height; i++) {
            for (int j = 0; j < width; j++) {
                double[] pixel = ballColor.get(i, j);
                blueValues[index] = pixel[0];
                greenValues[index] = pixel[1];
                redValues[index] = pixel[2];
                index++;
            }
        }

        // 计算每个通道的中位数
        double medianBlue = calculateMedian(blueValues);
        double medianGreen = calculateMedian(greenValues);
        double medianRed = calculateMedian(redValues);

        // 返回中位数颜色
        return new Scalar(medianBlue, medianGreen, medianRed);
    }

    private double calculateMedian(double[] values) {
        // 复制数组以避免修改原始数组
        double[] sortedValues = Arrays.copyOf(values, values.length);
        Arrays.sort(sortedValues);

        int length = sortedValues.length;
        if (length % 2 == 0) {
            // 偶数个元素，取中间两个元素的平均值
            return (sortedValues[length / 2 - 1] + sortedValues[length / 2]) / 2.0;
        } else {
            // 奇数个元素，取中间元素
            return sortedValues[length / 2];
        }
    }

    private TennisPoint filterByColorAndArea(TennisPoint p) {
        double area = p.getWidth() * p.getHeight();
        // 根据网球的大小和位置调整追踪精度
        if (area >= 3 && area <= 625) { // 假设网球的面积在3到25*25像素之间
            if (p.getMedianColor().val[2] >= 180 && p.getMedianColor().val[1] >= 190 && p.getMedianColor().val[0] < 180) {
                // RG通道高 B通道低
                return p;
            }
            if (p.getMedianColor().val[2] >= 150 && p.getMedianColor().val[2] < 200
                    && p.getMedianColor().val[1] >= 210 && p.getMedianColor().val[0] < 100) {
                // 偏黄色
                return p;
            } else if (p.getMedianColor().val[2] >= 200 && p.getMedianColor().val[1] >= 200 && p.getMedianColor().val[0] >= 180
                    && p.getCircularity() > 0.75 && p.getWidth() >= 3 && p.getWidth() <= 8) {
                // 偏白色，很小，且是一个近正圆
                return p;
            } else if (p.getMedianColor().val[0] <= 60
                    && p.getMedianColor().val[1] / (p.getMedianColor().val[0] == 0 ? 1 : p.getMedianColor().val[0]) >= 1.5
                    && p.getMedianColor().val[2] / (p.getMedianColor().val[0] == 0 ? 1 : p.getMedianColor().val[0]) >= 1.4) {
                // B通道特别低 绿色/蓝色>=1.6，红色/蓝色>=1.4， 蓝色<60
                return p;
            } else if (p.getMedianColor().val[1] / (p.getMedianColor().val[0] == 0 ? 1 : p.getMedianColor().val[0]) >= 1.45
                    && p.getMedianColor().val[2] / (p.getMedianColor().val[0] == 0 ? 1 : p.getMedianColor().val[0]) >= 1.35
                    && p.getWidth() >= 3 && p.getWidth() <= 16.5 && p.getCircularity() > 0.7) {
                // 远处球比较小的时候，光线比较暗 或者 近处光线暗
                return p;
            } else if (p.getMedianColor().val[1] > 200
                    && p.getMedianColor().val[1] / (p.getMedianColor().val[0] == 0 ? 1 : p.getMedianColor().val[0]) >= 1.6
                    && p.getMedianColor().val[2] / (p.getMedianColor().val[0] == 0 ? 1 : p.getMedianColor().val[0]) >= 1.1
                    && p.getWidth() >= 3 && p.getWidth() <= 16.5 && p.getCircularity() > 0.75) {
                // G>200时 但是R比较小 但是比B大，且圆度高 对于蓝色地胶比较有效果
                return p;
            } else {
                if (p.getCircularity() >= 0.80 && p.getWidth() >= 3 && p.getWidth() <= 20) {
                    // 其他情况 圆度高 且 宽度在3到20之间
                    return p;
                }
                p.setFilterType(2);
            }
        } else {
            p.setFilterType(1);
        }
        return p;
    }

}
